package com.me.study.concurrent.queue;

import com.me.common.pojo.Student;
import org.junit.Test;

import java.util.Comparator;
import java.util.PriorityQueue;

/**
 * PriorityQueue：优先级队列测试类

 * @author ME
 */
public class PriorityQueueTest {

    /**
     * 顶堆解决 Top N 问题：由于顶堆只保证最大、最小是有序的，因此 取最大 用 小顶堆，取最小 用 大顶堆。
     */
    @Test
    public void topN() {
        // 此处演示取最大4个。
        SmallestHeap<Student> heap = new SmallestHeap<>(4, Comparator.comparingInt(Student::getAge));
        heap.offer(new Student("Jack", 18));
        heap.offer(new Student("Tom", 14));
        heap.offer(new Student("Lucy", 19));
        heap.offer(new Student("Ljy", 17));
        System.out.println(heap);

        heap.offer(new Student("Page", 23));
        System.out.println(heap);
        heap.offer(new Student("Jimmy", 13));
        System.out.println(heap);
    }

    /**
     * PriorityQueue 实现容量固定的大顶堆
     */
    @Test
    public void maxHeap() {
        PriorityQueue<Student> queue = new PriorityQueue<>(4, (s1, s2) -> s2.getAge() - s1.getAge());
        queue.offer(new Student("Jack", 18));
        queue.offer(new Student("Tom", 14));
        queue.offer(new Student("Lucy", 19));
        queue.offer(new Student("Ljy", 17));
        System.out.println(queue);

        queue.poll();
        System.out.println(queue);
    }

    private static class SmallestHeap<T> extends PriorityQueue<T> {
        private final int maxSize;

        public SmallestHeap(int maxSize) {
            super(maxSize);
            this.maxSize = maxSize;
        }

        public SmallestHeap(int maxSize, Comparator<? super T> comparator) {
            super(maxSize, comparator);
            this.maxSize = maxSize;
        }

        @Override
        public boolean offer(T t) {
            // 如果未达到最大容量，则直接插入
            if (super.size() < maxSize) {
                return super.offer(t);
            }
            // 如果已经达到最大容量，并且小顶堆的堆顶元素 小于 新插入的值，则移除堆顶再插入
            else if (super.comparator().compare(super.peek(), t) < 0) {
                super.poll();
                return super.offer(t);
            }
            // 如果已经达到最大容量，并且小顶堆的堆顶元素 大于等于 新插入的值，则移除堆顶再插入
            else {
                return false;
            }
        }

    }

}
